import CR_DSPPytorch
import matplotlib.pyplot as plt
import numpy as np
import os
BASE_DIR = os.path.dirname(__file__)

if __name__ == "__main__":
    rect = CR_DSPPytorch.Rect(200, 2)
    rrc = CR_DSPPytorch.RRC(0.1, 10, 4)
    figure_dir = os.path.join(BASE_DIR, 'show_figure')
    if not os.path.exists(figure_dir):
        os.makedirs(figure_dir)

    rect_f = np.fft.fft(rect)

    plt.plot(rect, '-*')
    plt.savefig(os.path.join(figure_dir, 'rect_win_fd.png'))
    plt.close()

    plt.plot(rrc, '-*')
    plt.savefig(os.path.join(figure_dir, 'rrc_win_fd.png'))
    plt.close()

    plt.plot(rect_f.real, '-')
    plt.savefig(os.path.join(figure_dir, 'rect_win_fft.png'))
    plt.close()
